from solvers.alpha import Individual
import numpy as np
from typing import List


class Individual_3(Individual):
    def max_two(self, A_k: np.ndarray) -> np.ndarray:
        """
        找到最大和第二大的行元素
        """

        def _max_row(A: np.ndarray):
            assert A.shape.__len__() == 2
            array_blank = np.zeros_like(A)
            array_blank[range(A.shape[1]), np.argmax(np.abs(A), axis=1)] = A[
                range(A.shape[1]), np.argmax(np.abs(A), axis=1)
            ]
            return array_blank

        max_1_mat = _max_row(A_k)
        max_2_mat = _max_row(A_k - max_1_mat)
        return max_1_mat + max_2_mat

    def __call__(self, val: float = 1, sigma: float = 0.2) -> Individual:
        return Individual_3(val, sigma)

    def score(
        self, dft: np.ndarray, A_K: List[np.ndarray], K: int, beta: float
    ) -> float:
        if self._score != -np.inf:
            return self._score

        def _quantify(A: np.ndarray):
            B = np.zeros_like(A, dtype=int)
            A_real = np.real(A)
            A_imag = np.imag(A)
            B = (
                np.sign(A_real)
                * 2 ** np.ceil((np.clip(np.log2(np.abs(A_real)), 0, None)) - 1)
                + np.sign(A_imag)
                * 2 ** np.ceil((np.clip(np.log2(np.abs(A_imag)), 0, None)) - 1)
                * 1j
            )
            return B

        self._score = (
            -np.sqrt(
                np.sum(
                    np.abs(
                        beta * dft
                        - self.matmul(
                            map(lambda x: self.max_two(_quantify(
                                self.val * x)) / self.val, A_K), K
                        )
                    )
                )
            )
            / A_K[0].shape[0]
        )
        return self._score
